import os
import csv
import shutil
from pathlib import Path
from typing import Dict, Optional


SOURCE_AUDIO_DIR = "/hdd/data/swp/MER2025-process/audio"
CSV_PATH = "/hdd/data/swp/MER2025-process/track1_test_dis.csv"
DEST_BASE = "/hdd/data/swp/DH/MER2025/audio"


DISCRETE_TO_CATEGORY: Dict[str, str] = {
    # negative
    "worried": "negative",
    "angry": "negative",
    "sad": "negative",
    # neutral
    "neutral": "neutral",
    # positive
    "happy": "positive",
    "surprise": "positive",
}


def ensure_dest_dirs(base_dir: str) -> None:
    for sub in {"negative", "neutral", "positive"}:
        Path(base_dir, sub).mkdir(parents=True, exist_ok=True)


def load_name_to_discrete(csv_file: str) -> Dict[str, str]:
    mapping: Dict[str, str] = {}
    with open(csv_file, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        if "name" not in reader.fieldnames or "discrete" not in reader.fieldnames:
            raise ValueError("CSV must contain 'name' and 'discrete' columns")
        for row in reader:
            name_raw = (row.get("name") or "").strip()
            discrete_raw = (row.get("discrete") or "").strip().lower()
            if not name_raw or not discrete_raw:
                continue
            # store both with and without extension for robust matching
            stem = os.path.splitext(name_raw)[0]
            mapping[name_raw] = discrete_raw
            mapping[stem] = discrete_raw
    return mapping


def resolve_category(discrete: str) -> Optional[str]:
    return DISCRETE_TO_CATEGORY.get(discrete)


def copy_audio_by_discrete(source_dir: str, csv_path: str, dest_base: str) -> None:
    ensure_dest_dirs(dest_base)

    src_path = Path(source_dir)
    if not src_path.exists():
        raise FileNotFoundError(f"Source audio directory not found: {source_dir}")
    if not Path(csv_path).exists():
        raise FileNotFoundError(f"CSV file not found: {csv_path}")

    name_to_discrete = load_name_to_discrete(csv_path)

    total = 0
    matched = 0
    copied = 0
    skipped = 0

    for audio_path in src_path.rglob("*"):
        if not audio_path.is_file():
            continue
        total += 1
        name = audio_path.name
        stem = audio_path.stem

        discrete = name_to_discrete.get(name) or name_to_discrete.get(stem)
        if not discrete:
            skipped += 1
            print(f"[SKIP] Not listed in CSV: {audio_path}")
            continue

        matched += 1
        category = resolve_category(discrete)
        if not category:
            skipped += 1
            print(f"[SKIP] Unknown discrete '{discrete}' for file: {audio_path}")
            continue

        dest_dir = Path(dest_base, category)
        dest_file = dest_dir / name

        try:
            shutil.copy2(audio_path, dest_file)
            copied += 1
            print(f"[OK] {audio_path} -> {dest_file}")
        except Exception as e:
            skipped += 1
            print(f"[ERR] Failed to copy {audio_path} -> {dest_file}: {e}")

    print(
        f"\nDone. Total files scanned: {total}, Matched in CSV: {matched}, Copied: {copied}, Skipped: {skipped}.\n"
        f"Dest base: {dest_base}"
    )


if __name__ == "__main__":
    copy_audio_by_discrete(SOURCE_AUDIO_DIR, CSV_PATH, DEST_BASE)


